from ultralytics import YOLO
import cv2



# weight_path = r"E:\zhihao\2025文档\25年10月\ai_训练\权重\public_security\traffic_accient\20250925\train77\weights\best.pt"
weight_path = r"E:\zhihao\2025文档\25年10月\ai_训练\权重\public_security\traffic_accient\20251016\train89\weights\best.pt"
# video_path = r"E:\zhihao\2025文档\25年10月\ai_训练\视频图片_模型模拟\20251017\模拟模型\1013.mp4"
# video_path = r"E:\zhihao\2025文档\25年10月\ai_训练\视频图片_模型模拟\20251017\模拟模型\1016-1.mp4"
video_path = r"E:\zhihao\2025文档\25年10月\ai_训练\视频图片_模型模拟\20251017\模拟模型\1016-3.mp4"


# 打开视频文件
cap = cv2.VideoCapture(video_path,cv2.CAP_FFMPEG)
# 获取视频的基本信息
fps = int(cap.get(cv2.CAP_PROP_FPS))  # 获取原视频帧率
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))  # 视频宽度
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))  # 视频高度
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))  # 总帧数
print(fps,width,height,frame_count)

model = YOLO(weight_path)

# model.track(video_path, imgsz=(height,width),classes=[0], rect=True, iou=0.3, conf=0.6,visualize=False,vid_stride=2,half=True, save=False,show=False)
# model.predict(video_path,vid_stride=2,classes=[3,4,5], save=True,show=False,show_labels=True,show_conf=False)
#通用
model.predict(video_path,vid_stride=1,imgsz=(height,width),classes=[0],save=True,show=True,show_labels=True,show_conf=True,conf=0.5,iou=0.3)


